import numpy as np
from utils.config import *
from dataloader.data_loader import img_loader
from predict import predict
from save import save
from validate import validate

def main():
    print('dataset:', dataset)

    avg_error_avg, avg_error_1px, avg_error_2px, avg_error_3px, avg_error_4px, avg_error_5px = 0, 0, 0, 0, 0, 0

    for  batch in img_loader:
        # 从batch中取样本
        ground_truth = np.squeeze(batch['gt'])
        img_name = batch['name'][0]

        if dataset == 'sgbm': # 只做验证
            disp = np.squeeze(batch['sgbm'])
        else:
            left_img = np.squeeze(batch['left'])
            right_img = np.squeeze(batch['right'])

            disp = predict(left_img, right_img) # 生成视差
            save(disp, img_name) # 保存视差

        # 验证
        error = validate(disp, ground_truth, img_name)
        avg_error_avg += error['error_avg']
        avg_error_1px += error['error_1px']
        avg_error_2px += error['error_2px']
        avg_error_3px += error['error_3px']
        avg_error_4px += error['error_4px']
        avg_error_5px += error['error_5px']

    print('avg_error_avg:', avg_error_avg / len(img_loader))
    print('avg_error_1px:', avg_error_1px * 100 / len(img_loader))
    print('avg_error_2px:', avg_error_2px * 100 / len(img_loader))
    print('avg_error_3px:', avg_error_3px * 100 / len(img_loader))
    print('avg_error_4px:', avg_error_4px * 100 / len(img_loader))
    print('avg_error_5px:', avg_error_5px * 100 / len(img_loader))

if __name__ == '__main__':
    main()
